from __future__ import annotations
from ..utils.logger import get_logger
from collections import deque
from concurrent.futures import Executor
from weakref import WeakValueDictionary, WeakSet, ReferenceType, finalize  # noqa: F401
from typing import Any, List, TypeVar, Generic, Optional, Union, Dict
from dataclasses import field
import inspect
import weakref
import dataclasses
import types
import asyncio
import threading

T_in = TypeVar("T_in")
T_out = TypeVar("T_out")

# 用于标记无法被弱引用的类型
_NO_WEAKREF_BASIC_TYPES = (
    int,
    float,
    str,
    bool,  # bool 是 int 的子类，也是不可弱引用的
    types.NoneType,  # None 的类型
    tuple,
    bytes,
    list,
    dict,
    set,
    frozenset,
)

logger = get_logger("global")


@dataclasses.dataclass
class TraceWrapper(Generic[T_in, T_out]):
    data: T_in
    future: asyncio.Future[T_out] = field(
        default_factory=asyncio.Future
    )  # 用于存储异步结果的Future, 用于需要回流传递的任务
    sid: Optional[Union[str, int]] = field(default=None)  # (连接事务)哈希码
    chain: Dict[ReferenceType[TraceStream], float] = field(
        default_factory=dict
    )  # 用于记录数据来源的流模块和用时

    def __hash__(self):
        return id(
            self.data
        )  # 指向数据对象的内存地址, 保证唯一性, `update_data`后则会发生变化

    def __post_init__(self):
        if self.sid is None:
            # 如果未指定id, 则使用数据对象的内存地址作为id(当前唯一标识)
            self.sid = f"{id(self.data):x}"
        # finalize(
        #     self, lambda ident, chain: logger.debug(f">> {ident}, {chain}"), f"{abs(id(self)):016x}", self.chain
        # )

    def __repr__(self):
        return f"ID:{self.sid}|DATA:{type(self.data)}|CHAIN:{[m() for m in self.chain.keys()] if self.chain else 'UNKNOWN'}"

    def update_data(self, data: Union[T_in, T_out]):
        self.data = data
        return self

    def update_chain(self, stream: TraceStream, duration: float = 0.0):
        self.chain[weakref.ref(stream)] = duration
        return self

    @classmethod
    def copy_from(cls, wrapper: TraceWrapper) -> TraceWrapper:
        """
        data和future浅拷贝(rc), chain深拷贝(不同chain)
        适用于广播但是想要共用data和future的场景
        """
        return cls(
            data=wrapper.data,
            future=wrapper.future,
            sid=wrapper.sid,
            chain=wrapper.chain.copy(),
        )


class TraceStream(Generic[T_in, T_out]):
    """
    流模块abc, 继承后需要实现__call__, 并指定输入输出类型(静态类型检查)

    Class Attributes:
    -----------------
        is_broadcast (bool): 当多个下游时, 是否广播数据
        filter_input (bool): 是否检查输入类型(当上游广播时有效)
        run_in_background (bool, optional): 是否在后台运行call, 默认根据__call__返回值自动判断

    Args:
        type_in (T_in): 输入数据类型
        type_out (T_out): 输出数据类型
        executor (Executor, optional): 线程池. Defaults to None.
        max_queue_size (int, optional): 最大队列长度. Defaults to None.
        timeout_perflow (int, optional): 每次处理数据的超时时间, 超时则跳过. Defaults to 0.1.
    """

    # TraceStream的类属性, 相当于全局储存
    _traced_data: WeakValueDictionary[Union[str, int], TraceWrapper] = (
        WeakValueDictionary()
    )
    __slots__ = (
        "__weakref__",
        "_logger",
        "_loop_main",
        "_loop_porter",
        "_had_warning_no_consumer",
        "_executor",
        "is_coroutine",
        "signature",
        "need_future",
        "_type_in",
        "_type_out",
        "_up_streams",
        "_dn_streams",
        "_input_queue",
        "_buffer",
        "_is_broadcast",
        "_filter_input",
        "_buffers",
        "_is_working",
        "_start_event",
        "_timeout_perflow",
        "_running_count",
    )

    def __init__(
        self,
        *,
        executor: Executor = None,
        max_queue_size: int = None,
        timeout_perflow: int = 0.1,  # 缺省1秒
        is_broadcast: Optional[bool] = None,  # 当多个下游时, 是否广播数据
        filter_input: Optional[bool] = None,  # 是否检查输入类型(当上游广播时可以有效过滤入参数据)
    ):
        self._executor = executor
        
        # 使用从__init_subclass__解析的元数据
        self.signature = inspect.signature(self.__call__)
        self.is_coroutine = getattr(self.__class__, '_meta_is_coroutine', False)
        self.need_future = getattr(self.__class__, '_meta_need_future', False)
        resolved_type_in = getattr(self.__class__, '_meta_type_in', Any)
        resolved_type_out = getattr(self.__class__, '_meta_type_out', Any)

        self._type_in = resolved_type_in
        self._type_out = resolved_type_out
        self._up_streams: WeakSet[TraceWrapper[T_in, T_out]] = WeakSet()
        self._dn_streams: WeakSet[TraceWrapper[T_in, T_out]] = WeakSet()

        # 输入队列, 用于存放上游流传来的数据以及异步唤醒
        self._input_queue: asyncio.Queue[TraceWrapper[T_in, T_out]] = asyncio.Queue()
        self._buffer: deque[TraceWrapper[T_in, T_out]] = deque(
            maxlen=max_queue_size
        )  # 本Flow完成后的数据队列
        self._filter_input = filter_input if filter_input is not None else self._meta_filter_input
        self._is_broadcast = is_broadcast if is_broadcast is not None else self._meta_is_broadcast
        if self._is_broadcast:
            self._buffers: Dict[TraceStream, deque[TraceWrapper[T_in, T_out]]] = {}

        self._start_event: threading.Event = threading.Event()  # 启动信号
        self._timeout_perflow = timeout_perflow
        self._running_count = 0

    def __init_subclass__(
        cls, *, 
        is_broadcast: bool = False,  # 当多个下游时, 是否广播数据
        filter_input: bool = False,  # 是否检查输入类型(当上游广播时可以有效过滤入参数据),
        run_in_background: Optional[bool] = None,  # 是否在后台运行call (call返回为None时)
        **kwargs
    ):
        """
        类属性初始化
        is_broadcast: bool, 当多个下游时, 是否广播数据
        filter_input: bool, 是否检查输入类型(当上游广播时可以有效过滤入参数据)
        继承时指定, 也可以__init__初始化时指定
        
        自动解析__call__方法的类型注解和签名信息
        """
        super().__init_subclass__(**kwargs)
        cls._meta_is_broadcast = is_broadcast
        cls._meta_filter_input = filter_input
        
        # 解析__call__方法的签名信息
        signature = inspect.signature(cls.__call__)
        cls._meta_is_coroutine = inspect.iscoroutinefunction(cls.__call__)
        cls._meta_need_future = len(signature.parameters) == 3  # self + data + future
        
        # 解析__call__方法的类型注解
        params = list(signature.parameters.values())
        
        # 解析输入类型
        if len(params) >= 2:  # self + data [+ future]
            data_param = params[1]  # 跳过self参数
            if data_param.annotation != inspect._empty:
                cls._meta_type_in = data_param.annotation
            else:
                cls._meta_type_in = types.NoneType
        else:
            cls._meta_type_in = types.NoneType
            
        # 解析输出类型 (返回值注解)
        return_annotation = signature.return_annotation
        if return_annotation != inspect._empty:
            # 处理可能的泛型类型 (如 List[TraceWrapper])
            if hasattr(return_annotation, '__origin__'):
                if return_annotation.__origin__ is list:
                    # 如果返回List[T]，提取T作为输出类型
                    if hasattr(return_annotation, '__args__') and return_annotation.__args__:
                        cls._meta_type_out = return_annotation.__args__[0]
                    else:
                        cls._meta_type_out = return_annotation
                else:
                    cls._meta_type_out = return_annotation
            else:
                cls._meta_type_out = return_annotation
        else:
            cls._meta_type_out = types.NoneType

        # 将None转换为types.NoneType
        cls._meta_type_in = types.NoneType if cls._meta_type_in is None else cls._meta_type_in
        cls._meta_type_out = types.NoneType if cls._meta_type_out is None else cls._meta_type_out
            
        # 确定是否在后台运行
        if run_in_background is None:
            match return_annotation:
                case None | types.NoneType: # 明确指定返回None
                    cls._meta_run_in_background = True
                case inspect._empty: # 未指定返回值
                    cls._meta_run_in_background = False
                case _: # 其他返回值
                    cls._meta_run_in_background = False
        else:
            cls._meta_run_in_background = run_in_background

    @property
    def type_info(self) -> dict:
        """返回类型信息，用于调试"""
        return {
            'type_in': self._type_in,
            'type_out': self._type_out,
            'meta_type_in': getattr(self.__class__, '_meta_type_in', None),
            'meta_type_out': getattr(self.__class__, '_meta_type_out', None),
            'meta_is_coroutine': getattr(self.__class__, '_meta_is_coroutine', None),
            'meta_need_future': getattr(self.__class__, '_meta_need_future', None),
            'meta_run_in_background': getattr(self.__class__, '_meta_run_in_background', None),
            'signature': str(self.signature),
            'is_coroutine': self.is_coroutine,
            'need_future': self.need_future,
        }

    def __repr__(self):
        return f"{self.__class__.__name__}:{id(self):x}"

    @property
    def loop_main(self) -> Optional[asyncio.AbstractEventLoop]:
        loop = getattr(self, "_loop_main", None)
        return loop

    @loop_main.setter
    def loop_main(self, value: Optional[asyncio.AbstractEventLoop]):
        assert isinstance(value, Optional[asyncio.AbstractEventLoop]), (
            "loop must be an instance of asyncio.AbstractEventLoop"
        )
        self._loop_main = value

    @property
    def loop_porter(self) -> asyncio.AbstractEventLoop:
        loop = getattr(self, "_loop_porter", None)
        return loop

    @loop_porter.setter
    def loop_porter(self, value: asyncio.AbstractEventLoop):
        assert isinstance(value, asyncio.AbstractEventLoop), (
            "loop must be an instance of asyncio.AbstractEventLoop"
        )
        self._loop_porter = value

    @property
    def checkin(self):
        return (self._type_in, TraceWrapper, types.NoneType, type(...))

    @property
    def checkout(self):
        return (self._type_out, TraceWrapper, types.NoneType, type(...))

    @property
    def traced_data(self) -> WeakValueDictionary[Union[str, int], TraceWrapper]:
        """
        用于存储数据引用, 确保数据位置必定在缓冲区或__call__中
        弱引用保证一旦数据生命周期完结可以被gc
        """
        trace_data = getattr(self.__class__, "_traced_data", None)
        if trace_data is None:
            self.logger.error(
                "traced_data应该是类属性, 不该为None"
            )  # 理论上never reach
            self._trace_data: WeakValueDictionary[Union[str, int], TraceWrapper] = (
                WeakValueDictionary()
            )
            trace_data = self._trace_data
        return trace_data

    @property
    def logger(self):
        logger = getattr(self, "_logger", None)
        if logger is None:
            # 针对同类不同实例需要不同的logger
            self._logger = logger = get_logger(f"{self}")  # __repr__
        return logger

    @property
    def is_broadcast(self):
        return self._is_broadcast and self.dn_streams.__len__() > 0

    @property
    def up_streams(self) -> WeakSet[TraceStream]:
        """获取上游流"""
        if not hasattr(self, "_up_streams"):
            self._up_streams = WeakSet()
        return self._up_streams

    @property
    def dn_streams(self) -> WeakSet[TraceStream]:
        """获取下游流"""
        if not hasattr(self, "_dn_streams"):
            self._dn_streams = WeakSet()
        return self._dn_streams

    def hook_to(
        self,
        module: TraceStream,
    ) -> TraceStream:
        "设置上游, 不被反复触发"
        self.up_streams.add(module)
        return module

    def unhook(self, module: TraceStream):
        "解除上游, 不被反复触发"
        self.up_streams.discard(module)

    def link_to(
        self,
        module: TraceStream,
    ) -> TraceStream:
        "设置下游"
        self.dn_streams.add(module)
        module.hook_to(self)  # 设置双向链接
        if self.is_broadcast:
            self._buffers[module] = deque(maxlen=self._buffer.maxlen)
        return module

    def unlink(self, module: TraceStream):
        "解除下游"
        self.dn_streams.discard(module)
        if self.is_broadcast:
            buffer = self._buffers.pop(module)
            buffer.clear()
        module.unhook(self)  # 解除双向链接

    async def export(
        self,
        wrapper: Union[TraceWrapper, List[TraceWrapper]],
        include_list: bool = True,
    ):
        if self.dn_streams.__len__() == 0:
            if not hasattr(self, "_had_warning_no_consumer"):
                self.logger.warning(
                    "Has no downstream consumer, data will be remain, it is risky"
                )
                self._had_warning_no_consumer = True

        assert isinstance(wrapper, TraceWrapper), (
            f"Expected TraceWrapper, got {type(wrapper)}"
        )
        if wrapper is None:  # None默认不处理
            return
        if isinstance(wrapper, list) and include_list:
            for item in wrapper:
                await self.export(item, False)  # 谨防无限递归
            return
        self.trace(wrapper)  # 经手模块都记录该数据
        if self.is_broadcast:
            for buffer in self._buffers.values():
                buffer.append(TraceWrapper.copy_from(wrapper))
        else:
            self._buffer.append(wrapper)

    def input(
        self,
        data: Union[T_in, TraceWrapper],
    ):
        """
        外部调用put直接放入数据队列, 调试时可以使用
        """
        if isinstance(data, TraceWrapper):
            wrapper = data
        elif isinstance(data, self._type_in):
            wrapper = TraceWrapper(data, chain={weakref.ref(self): 0.0})
        elif data is None:
            # 停止信号
            wrapper = TraceWrapper(data, chain={weakref.ref(self): 0.0})
        else:
            raise TypeError(
                f"{self}: put data {data} is not a TraceWrapper or {self._type_out}"
            )
        self._input_queue.put_nowait(wrapper)  # 无等待放入中央队列

    def pull(self, puller: TraceStream):
        # caller_frame = inspect.currentframe().f_back
        # puller = caller_frame.f_locals.get("self")
        buffer = self._buffer
        if self.is_broadcast:
            if puller not in self._buffers:
                self._buffers[puller] = deque(self._buffer, maxlen=self._buffer.maxlen)
            buffer = self._buffers[puller]
        if buffer.__len__() == 0:
            return ...
        return buffer.popleft()

    async def run_porter(self, bg_loop: asyncio.AbstractEventLoop):
        """
        将上游数据拉入中央队列, 并且监听stop_event
        """
        self.loop_porter = bg_loop
        self._start_event.wait()
        while True:
            ups = [m for m in self.up_streams]
            if not ups:  # 最上游等待input即可
                await asyncio.sleep(0.1)
                return
            for up in ups:
                wrapper: TraceWrapper = up.pull(self)
                if wrapper == ...:
                    continue
                if self._filter_input and not isinstance(wrapper.data, self.checkin):
                    continue
                asyncio.run_coroutine_threadsafe(
                    self._input_queue.put(wrapper), self.loop_main
                )
            await asyncio.sleep(0.01)  # 10ms

    async def run_main(self, loop: asyncio.AbstractEventLoop):
        self._running_count += 1
        self.loop_main = loop
        self._start_event.set()
        while True:
            await self.flow()

    def stop_check(
        self, chain: Optional[dict[ReferenceType[TraceStream], float]] = None
    ):
        # 检查停止点
        # frame = inspect.currentframe().f_back
        # filename = frame.f_code.co_filename
        # lineno = frame.f_lineno
        # self.logger.debug(f"stop_check called from {frame.f_locals.get("self")} {filename}:{lineno}")
        self._running_count = max(0, self._running_count - 1)  # 减少运行计数
        # 通知传输链检查是否可以停止
        if chain is not None:
            for ref in chain.keys():
                if (stream := ref()) == self:
                    continue
                stream.stop_check()
        if self._running_count > 0:
            return
        self._start_event.clear()  # 清除启动信号
        raise StopAsyncIteration(self)

    async def flow(self) -> bool:
        """
        Returns:
            is_timeout: bool, 是否超时, True表示超时跳过, False表示正常处理
        """
        try:
            # 获取中央队列数据, 超时则略过(避免堵塞)
            wrapper_in: TraceWrapper = await asyncio.wait_for(
                self._input_queue.get(), timeout=self._timeout_perflow
            )
        except asyncio.TimeoutError:
            return True  # 超时则跳过
        if wrapper_in.data is None:
            # 收到停止信号透传(仅从顶层能有效发起, 需要input(None)或者TraceWrapper(None))
            if self.dn_streams.__len__() == 0:  # 无下游时
                self.stop_check(wrapper_in.chain)  # 唯一退出点
            else:
                await self.export(wrapper_in.update_chain(self, 0.0))  # 透传停止信号
            return False
        time = self.loop_main.time()  # 从获取开始计时
        if (
            self._type_in == TraceWrapper
        ):  # 如果底层要求是TraceWrapper或是其子类,则直接传入
            inputs = [wrapper_in]
        else:
            inputs = (
                [wrapper_in.data, wrapper_in.future]
                if self.need_future
                else [wrapper_in.data]
            )
        async def run_after_data():
            if self.is_coroutine:
                output = await self.__call__(*inputs)
            else:
                task = self.loop_main.run_in_executor(
                    self._executor, self.__call__, *inputs
                )
                output = await task
            if output is None:
                return False
            duration = self.loop_main.time() - time  # 计算处理时长
            if isinstance(output, list):
                await self.export(
                    [
                        item.update_chain(self, duration)
                        if isinstance(item, TraceWrapper)
                        else wrapper_in.update_data(item).update_chain(self, duration)
                        for item in output
                        if isinstance(item, self.checkout)
                    ]
                )
            elif isinstance(output, TraceWrapper):  # 自行封装
                output.update_chain(self, duration)  # 更新链路信息
                await self.export(output)
            elif isinstance(output, self._type_out):
                await self.export(
                    wrapper_in.update_data(output).update_chain(self, duration)
                )
            else:
                raise TypeError(f"{self}: Output type {type(output)} not in {self.checkout}")
            return False
        if self._meta_run_in_background:
            asyncio.create_task(run_after_data())
            return False
        else:
            return await run_after_data()

    def link(self, stream: TraceStream):
        from .pipeline import BasePipeline  # 避免循环导入

        self.link_to(stream)
        return BasePipeline(self, stream)

    def trace(self, wrapper: TraceWrapper):
        """
        记录数据流向, 用于调试, 必须弱引用
        """
        if wrapper.data is not None:
            self.traced_data[wrapper.sid] = wrapper

    def __call__(self, data: T_in, future: Optional[asyncio.Future[T_out]]) -> T_out:
        """
        Args:
            data (T_in): 输入数据
            future (Optional[asyncio.Future[T_out]]): 用于存储异步结果的Future, 用于需要回流传递的任务

        Returns:
            out (T_out): 输出数据

        需要在此处理每次进来的数据如何处理，
        然后放进本模块的队列中
        可以用async def实现, 将会放置在main loop执行
        入参签名可以使用data+future或者单独data, 
        对入参类型必须使用静态类型检查才能继承元类
        """
        assert self._type_in == self._type_out, (
            f"Cannot return data from {self._type_in} to {self._type_out} directly, Impelement __call__ to process data!"
        )
        return data  # 默认直接返回数据, 需要重写该方法来处理数据
